#!/usr/bin/python2.6
# -*- coding: utf-8 -*-
# Copyright (c) 2008 Qtrac Ltd. All rights reserved.
# This program or module is free software: you can redistribute it and/or
# modify it under the terms of the GNU General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version. It is provided for educational
# purposes and is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# General Public License for more details.


import sys
import pickle
import string
import os
import io
import numpy as np
from scipy import stats
import time
import CaseConAlgo
#
#def epiCaseControlOOInit():
#
#	#Build jobs
#	jobCC = {}
#	jobCC["Project"]="To be inserted from the GUI"#TODO from GUI
#	jobCC["Pklfiles"] = ["epi_results1.pkl","epi_results6.pkl","epi_results10.pkl","epi_results14.pkl"]#TODO selected from GUI
#
#	default_significance_level = 0.05#TODO from GUI
#	jobCC["NomSignLevel"]=default_significance_level
#	jobCC["TraitSearch"]=jobCC["Pklfiles"][0]
#			
#	#Listing traits TODO should pop up in GUI
#	trait_list = []
#	prime_res = open(jobCC["TraitSearch"],"rb")
#	prime_data = pickle.load(prime_res)
#	for trait in prime_data["epistasis_data"]:#trait valg? GUI? 
#		trait_select = trait["trait_name"]
#		trait_list.append(trait_select)
#	jobCC["TraitSel"]=trait_list
#
#	epiCaseControlOOTr(jobCC)



def start(class_pkl_files, significance_level, trait_list, project_name, analysis_root_dir, callback=None):

	#Build jobs
	jobCC = {}
	jobCC["Project"]=project_name
	jobCC["Pklfiles"] = class_pkl_files 
	jobCC["NomSignLevel"]= significance_level 
	jobCC["TraitSel"]=trait_list

	epiCaseControlOOTr(jobCC, analysis_root_dir)

    #jobCC["Pklfiles"] = ["epi_results1.pkl","epi_results6.pkl","epi_results10.pkl","epi_results14.pkl"]#TODO selected from GUI
    # jobCC["TraitSearch"]=jobCC["Pklfiles"][0]
#            
#    #Listing traits TODO should pop up in GUI
#    trait_list = []
#    prime_res = open(jobCC["TraitSearch"],"rb")
#    prime_data = pickle.load(prime_res)
#    for trait in prime_data["epistasis_data"]:#trait valg? GUI? 
#        trait_select = trait["trait_name"]
#        trait_list.append(trait_select)
    

def epiCaseControlOOTr(jobCC, analysis_root_dir):

#establish classes to test from pkl files

	for tr in jobCC["TraitSel"]:#replace with GUI choises of traits, a job is defined!
		new_summary_dir = os.path.join(analysis_root_dir, "Case-control One by One "+ tr +"/")
		if not os.path.isdir(new_summary_dir):#create output dir
			os.mkdir(new_summary_dir)	
		jobCC["OutputDir"]=new_summary_dir
		
		class_content = []
		class_numb = []
		class_name = []	
		for epi in jobCC["Pklfiles"]:#selected pkl-files from GUI
			all_res = open(epi,"rb")
			all_data = pickle.load(all_res)
			class_select = all_data["class"]
			
			for trait in all_data["epistasis_data"]:#trait valg?  
				trait_select = trait["trait_name"]
				if tr == trait_select:
					class_numb.append(class_select)				
					al_fil = "GAfile" + str(class_select)
					class_name.append(al_fil)
					al_fil = trait["AllGenFreq"]#A list of genotype and allel frequences as strings
					al_fil_mod = []
					for ga in range(0,len(al_fil)):
						al_li = al_fil[ga].split()
						al_fil_mod.append(al_li)
					class_content.append(al_fil_mod)
					
		jobCC["ClassListName"]=class_name
		jobCC["ClassListNo"]=class_numb		
		jobCC["ClassListData"]=class_content

		case_control_data(jobCC)

def case_control_data(jobCC):

	#Select epifiles. This should be a selection of all epi-files two at the time.
	#
	"""First epifile"""
	for epfil1 in range(0,len(jobCC["ClassListNo"])-1):	
		#GAfile number, used for ordering output
		jobCC["File1"] = jobCC["ClassListName"][epfil1][6:]		
		epifil_number1=jobCC["ClassListData"][epfil1]
		jobCC["Raw1"] = epifil_number1
		snplist1 = CaseConAlgo.snp_list_create(epifil_number1)

		"""Second epifile"""
		for epfil2 in range(epfil1+1,len(jobCC["ClassListNo"])):	
			jobCC["File2"] = jobCC["ClassListName"][epfil2][6:]
			epifil_number2=jobCC["ClassListData"][epfil2]
			jobCC["Raw2"] = epifil_number2		
			snplist2 = CaseConAlgo.snp_list_create(epifil_number2)

			#Only snps present both epifiles are of interest
			snp_list_common = list(set(snplist1[0]) & set(snplist2[0]))
			jobCC["SnpCommon"]=snp_list_common

			#Significance levels based on number of snps-tested
			sign_level_nominal= jobCC["NomSignLevel"]#from GUI
			numb_tests = len(snp_list_common)*(len(snp_list_common)-1)/2
			if numb_tests>0:
				sign_level_bonf = sign_level_nominal/numb_tests
			else:
				sign_level_bonf = "No common SNPs"


#			sign_level_bonf = sign_level_nominal/numb_tests
			jobCC["SignNominal"]=sign_level_nominal	
			jobCC["SignBonf"]=sign_level_bonf
			
			#Create output files
			create_output_files(jobCC)

			#Start calculations
			cascon_select_data(jobCC)
			
def create_output_files(jobCC):
	#Just to make output ordered
	first_epi=min(int(jobCC["File1"]),int(jobCC["File2"]))
	second_epi=max(int(jobCC["File1"]),int(jobCC["File2"]))
	jobCC["EpiA"]=str(first_epi)
	jobCC["EpiB"]=str(second_epi)

	#Create output files
	#Output nominal
	snps_summary_00N = open(jobCC["OutputDir"]+"/"+"GAfiles "+jobCC["EpiA"]+ " vs "+jobCC["EpiB"]+" Nominal", "w")
	jobCC["SumAllListN"]=snps_summary_00N	
	summ001 = time.asctime()+"\n\n"
	summ002 = "All SNPs compared pair-wise between subpopulations"+"\n\n"
	summ002a = "Subpopulations compared: "+str(jobCC["EpiA"])+  " and " + str(jobCC["EpiB"])+"\n\n"
	summ002b = "Number of SNPs compared: "+str(len(jobCC["SnpCommon"]))+ "\n\n"
	jobCC["SumAllListN"].write(summ001)
	jobCC["SumAllListN"].write(summ002)
	jobCC["SumAllListN"].write(summ002a)	
	jobCC["SumAllListN"].write(summ002b)
	first_class = "Subpopulation "+str(jobCC["EpiA"])
	second_class = "Subpopulation "+str(jobCC["EpiB"])
	summ003 = "Significance level (Nominal): "+str(jobCC["SignNominal"])+"\n\n"
	summ004 = "{0} {1} {2} {3} {4} {5}\n".format("\t","\t"+first_class+"\t","\t","\t","\t",second_class)	
	summ004a = "{0} {1} {2} {3} {4} {5} {6} {7} {8} {9} {10}\n\n".format("SNP"+"\t","Gene"+"\t","AA"+"\t","Aa"+"\t","aa"+"\t","\t","BB"+"\t","Bb"+"\t","bb"+"\t","Stat_value"+"\t","p_value"+"\t")

#	summ004 = "{0} {1} {2} {3} {4} {5} {6} {7} {8} {9} {10} {11} \n\n".format("SNP","Gene","EpifileA","AA","Aa","aa","EpifileB","BB","Bb","bb","Stat_value","p_value")
	jobCC["SumAllListN"].write(summ003)
	jobCC["SumAllListN"].write(summ004)
	jobCC["SumAllListN"].write(summ004a)
	
	#Output Bonferoni
	snps_summary_00B = open(jobCC["OutputDir"]+"/"+"GAfiles "+jobCC["EpiA"]+ " vs "+jobCC["EpiB"]+" Bonferoni", "w")
	jobCC["SumAllListB"]=snps_summary_00B	
	summ001 = time.asctime()+"\n\n"
	summ002 = "All SNPs compared pair-wise between subpopulations"+"\n\n"
	jobCC["SumAllListB"].write(summ001)
	jobCC["SumAllListB"].write(summ002)
	jobCC["SumAllListB"].write(summ002a)
	jobCC["SumAllListB"].write(summ002b)
	summ005 = "Significance level (Bonferoni): "+str(jobCC["SignBonf"])+"\n\n"
	jobCC["SumAllListB"].write(summ005)
	jobCC["SumAllListB"].write(summ004)
	jobCC["SumAllListB"].write(summ004a)
				
def cascon_select_data(jobCC):#own script

	#Create two files from each files selected containing the selected snp
	for episnp in range(0,len(jobCC["SnpCommon"])):
		snp_select=jobCC["SnpCommon"][episnp]
		jobCC["SnpSelect"]=snp_select
		#Extract entries with selected snp from the first file
		data_raw1_select=[]
		for epis in range(0,len(jobCC["Raw1"])):
			if jobCC["Raw1"][epis][0] == snp_select:
				data_raw1_select.append(jobCC["Raw1"][epis])
				genename1 = jobCC["Raw1"][epis][1]
			if jobCC["Raw1"][epis][7] == snp_select:
				data_raw1_select.append(jobCC["Raw1"][epis])
				genename1 = jobCC["Raw1"][epis][8]
		jobCC["DataCompile1"]=data_raw1_select
		jobCC["GeneName1"]=genename1

		#Extract entries with selected snp from the second file
		data_raw2_select=[]
		for epis in range(0,len(jobCC["Raw2"])):
			if jobCC["Raw2"][epis][0] == snp_select:
				data_raw2_select.append(jobCC["Raw2"][epis])
				genename2 = jobCC["Raw2"][epis][1]
			if jobCC["Raw2"][epis][7] == snp_select:
				data_raw2_select.append(jobCC["Raw2"][epis])
				genename2 = jobCC["Raw2"][epis][8]

		jobCC["DataCompile2"]=data_raw2_select
		
		#Call script to calulate case-control
		summ_list = cascon_select_calc(jobCC)

#		print summ_list

		if summ_list[1]<jobCC["SignNominal"]:
#			snps_sign_head2 = "{0} {1} {2} {3} {4} {5} {6} {7} {8} {9} {10}\n".format(jobCC["SnpSelect"],jobCC["GeneName1"],jobCC["GenoList1"][0],jobCC["GenoList1"][1],jobCC["GenoList1"][2],"\t",jobCC["GenoList2"][0],jobCC["GenoList2"][1],jobCC["GenoList2"][2],summ_list[0],summ_list[1])
			snps_sign_head2 = "{0} {1} {2} {3} {4} {5} {6} {7} {8} {9} {10}\n".format(str(jobCC["SnpSelect"])+"\t",str(jobCC["GeneName1"])+"\t",str(jobCC["GenoList1"][0])+"\t",str(jobCC["GenoList1"][1])+"\t",str(jobCC["GenoList1"][2])+"\t","\t",str(jobCC["GenoList2"][0])+"\t",str(jobCC["GenoList2"][1])+"\t",str(jobCC["GenoList2"][2])+"\t",str(summ_list[0])+"\t",str(summ_list[1])+"\t")
			
			jobCC["SumAllListN"].write(snps_sign_head2)
			
		#Write significant result only for snp_select to summary files
		if summ_list[1] < jobCC["SignBonf"]:
			snps_sign_head2 = "{0} {1} {2} {3} {4} {5} {6} {7} {8} {9} {10}\n".format(str(jobCC["SnpSelect"])+"\t",str(jobCC["GeneName1"])+"\t",str(jobCC["GenoList1"][0])+"\t",str(jobCC["GenoList1"][1])+"\t",str(jobCC["GenoList1"][2])+"\t","\t",str(jobCC["GenoList2"][0])+"\t",str(jobCC["GenoList2"][1])+"\t",str(jobCC["GenoList2"][2])+"\t",str(summ_list[0])+"\t",str(summ_list[1])+"\t")
		
#			snps_sign_head3 = "{0} {1} {2} {3} {4} {5} {6} {7} {8} {9} {10}\n".format(jobCC["SnpSelect"],jobCC["GeneName1"],jobCC["GenoList1"][0],jobCC["GenoList1"][1],jobCC["GenoList1"][2],"\t",jobCC["GenoList2"][0],jobCC["GenoList2"][1],jobCC["GenoList2"][2],summ_list[0],summ_list[1])
			jobCC["SumAllListB"].write(snps_sign_head2)
		
def cascon_select_calc(jobCC):

	#Subjobs are created

	"""Calcualtions are performed using the largest numbers of subjects for each SNP. There may be some minor differences as the allelfrequency files are enumerated for two-SNP interactions when calculating two-SNP epistasis. However, these minor differences are not of interest; rather the larger the number of subjects the power of the case-control calculations will increase. This is justified, as lower numbers just mean that the genotype has beee filtered, but all filtered genotypes are actually present (and maybe more) in the basic suppopulation.
#Therfore consider using the basic populations, not those filtered by epistasis calculations"""

	#List genotypes from first subpopulation (epifile). Necessary as the selected SNP may in column 0 or 7

	geno_types1 = []
	#From column 0..	
	for gt in range(0,len(jobCC["DataCompile1"])):
		if jobCC["DataCompile1"][gt][0] == jobCC["SnpSelect"]:
			geno_types1.append(jobCC["DataCompile1"][gt][2:5])#strange, ony the first 3 to be used
	#From column 7..	
	for gt in range(0,len(jobCC["DataCompile1"])):
		if jobCC["DataCompile1"][gt][7] == jobCC["SnpSelect"]:
			geno_types1.append(jobCC["DataCompile1"][gt][9:12])

	genot1_list = []
	max_numb = 0	
	for gnt in range(0,len(geno_types1)):
		gg= geno_types1[gnt]
		gt_num = int(gg[0])+int(gg[1])+int(gg[2])		
		if gt_num>max_numb:
			max_numb = gt_num
			genot1_list = [int(gg[0]),int(gg[1]),int(gg[2])]		
	jobCC["GenoList1"]=genot1_list

	#List from second subpopulation
	geno_types2 = []
	#From column 0..	
	for gt in range(0,len(jobCC["DataCompile2"])):
		if jobCC["DataCompile2"][gt][0] == jobCC["SnpSelect"]:
			geno_types2.append(jobCC["DataCompile2"][gt][2:5])
	#From column 7..	
	for gt in range(0,len(jobCC["DataCompile2"])):
		if jobCC["DataCompile2"][gt][7] == jobCC["SnpSelect"]:
			geno_types2.append(jobCC["DataCompile2"][gt][9:12])

	#Find max in second subpopulation 
	genot2_list = []
	max_numb = 0		
	for gnt in range(0,len(geno_types2)):
		gg= geno_types2[gnt]
		gt_num = int(gg[0])+int(gg[1])+int(gg[2])	
		if gt_num>max_numb:
			max_numb = gt_num
			genot2_list = [int(gg[0]),int(gg[1]),int(gg[2])]		
	jobCC["GenoList2"]=genot2_list

	#Find most significant values

	calcepi=CaseConAlgo.cascon_calc_proper(jobCC)

	return calcepi


if __name__ == '__main__':
    start(class_pkl_files=["epi_results1.pkl", "epi_results2.pkl"], significance_level=0.05, trait_list=["ApoB"], project_name="test,", analysis_root_dir=".")

#    print("1) Script 'calc_effects' is called")

#epiCaseControlOOInit()


